// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use rand::{rngs::ReseedingRng, RngCore};
use rand_chacha::ChaChaCore;

pub use s2n_quic_core::random::*;

struct AwsLc;

impl RngCore for AwsLc {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        let mut v = [0; 4];
        self.fill_bytes(&mut v);
        u32::from_ne_bytes(v)
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        let mut v = [0; 8];
        self.fill_bytes(&mut v);
        u64::from_ne_bytes(v)
    }

    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        aws_lc_rs::rand::fill(dest).unwrap()
    }
}

pub struct Random {
    public: ReseedingRng<ChaChaCore, AwsLc>,
    private: ReseedingRng<ChaChaCore, AwsLc>,
}

impl Default for Random {
    #[inline]
    fn default() -> Self {
        Self {
            public: build_rng(),
            private: build_rng(),
        }
    }
}

// Constructs a `ReseedingRng` with a ChaCha RNG initially seeded from the OS,
// that will reseed from the OS after RESEED_THRESHOLD is exceeded
fn build_rng() -> ReseedingRng<ChaChaCore, AwsLc> {
    // Number of generated bytes after which to reseed the public and private random
    // generators.
    //
    // This value is based on THREAD_RNG_RESEED_THRESHOLD from
    // [rand::rngs::thread.rs](https://github.com/rust-random/rand/blob/ef75e56cf5824d33c55622bf84a70ec6e22761ba/src/rngs/thread.rs#L39)
    const RESEED_THRESHOLD: u64 = 1024 * 64;
    ReseedingRng::<ChaChaCore, AwsLc>::new(RESEED_THRESHOLD, AwsLc)
        .unwrap_or_else(|err| panic!("could not initialize random generator: {err}"))
}

impl Generator for Random {
    #[inline]
    fn public_random_fill(&mut self, dest: &mut [u8]) {
        self.public.fill_bytes(dest);
    }

    #[inline]
    fn private_random_fill(&mut self, dest: &mut [u8]) {
        self.private.fill_bytes(dest);
    }
}
